import sys
import struct
import serial
import csv
from datetime import datetime
from collections import namedtuple, defaultdict
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, 
                             QWidget, QLabel, QComboBox, QPushButton, QSpinBox, 
                             QFileDialog, QMessageBox)
from PyQt5.QtCore import QTimer, Qt
import pyqtgraph as pg
from serial.tools import list_ports

# 定义数据结构
IMUData = namedtuple('IMUData', ['pitch', 'roll', 'yaw', 'temperature'])
FrameData = namedtuple('FrameData', ['imu1', 'imu2', 'imu3', 'imu4', 'data1', 'data2', 'data3', 'data4'])

class IMUParser:
    @staticmethod
    def parse_imu_data(data_bytes):
        """解析单个IMU的数据（12字节浮点数 + 2字节温度）"""
        pitch, roll, yaw = struct.unpack('<fff', data_bytes[:12])  # 小端序
        temperature = struct.unpack('<H', data_bytes[12:14])[0]    # 小端序无符号短整型
        return IMUData(pitch, roll, yaw, temperature)

    @staticmethod
    def parse_frame(data_bytes):
        """解析一帧数据（66字节）"""
        # 检查帧头
        header = data_bytes[:2]
        if header != b'\xA5\x5A':
            raise ValueError(f"无效的帧头: {header.hex()}")
        
        # 解析4个IMU数据
        imu1 = IMUParser.parse_imu_data(data_bytes[2:16])
        imu2 = IMUParser.parse_imu_data(data_bytes[16:30])
        imu3 = IMUParser.parse_imu_data(data_bytes[30:44])
        imu4 = IMUParser.parse_imu_data(data_bytes[44:58])
        
        # 解析4个16位数据
        data1, data2, data3, data4 = struct.unpack('<4h', data_bytes[58:66])    # 小端序
        
        return FrameData(imu1, imu2, imu3, imu4, data1, data2, data3, data4)

class DataSaver:
    def __init__(self):
        self.records = []
        self.fieldnames = ['timestamp']
        
        # 添加IMU数据字段
        for i in range(1, 5):
            self.fieldnames.extend([
                f'imu{i}_pitch', f'imu{i}_roll', f'imu{i}_yaw', f'imu{i}_temp'
            ])
        
        # 添加附加数据字段
        self.fieldnames.extend(['data1', 'data2', 'data3', 'data4'])
    
    def add_record(self, frame_data):
        """添加一帧数据记录"""
        record = {'timestamp': datetime.now().isoformat(timespec='milliseconds')}
        
        # 添加IMU数据
        for i, imu in enumerate([frame_data.imu1, frame_data.imu2, frame_data.imu3, frame_data.imu4], 1):
            record.update({
                f'imu{i}_pitch': imu.pitch,
                f'imu{i}_roll': imu.roll,
                f'imu{i}_yaw': imu.yaw,
                f'imu{i}_temp': imu.temperature
            })
        
        # 添加附加数据
        record.update({
            'data1': frame_data.data1,
            'data2': frame_data.data2,
            'data3': frame_data.data3,
            'data4': frame_data.data4
        })
        
        self.records.append(record)
    
    def save_to_csv(self, filename):
        """保存数据到CSV文件"""
        if not self.records:
            return False
        
        try:
            with open(filename, 'w', newline='') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=self.fieldnames)
                writer.writeheader()
                writer.writerows(self.records)
            return True
        except Exception as e:
            print(f"保存文件时出错: {e}")
            return False
    
    def clear(self):
        """清空记录"""
        self.records.clear()

class MultiPlotWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.init_ui()
        self.data_history = defaultdict(list)
        self.max_points = 500  # 显示的最大数据点数
        
    def init_ui(self):
        self.layout = QVBoxLayout(self)
        
        # 创建图形布局
        self.graphics_layout = pg.GraphicsLayoutWidget()
        self.graphics_layout.setBackground('w')
        self.layout.addWidget(self.graphics_layout)
        
        # 创建4个绘图区域
        self.plot_pitch = self.graphics_layout.addPlot(title="Pitch角度 (°)", row=0, col=0)
        self.plot_roll = self.graphics_layout.addPlot(title="Roll角度 (°)", row=1, col=0)
        self.plot_yaw = self.graphics_layout.addPlot(title="Yaw角度 (°)", row=0, col=1)
        self.plot_data = self.graphics_layout.addPlot(title="附加数据", row=1, col=1)
        
        # 设置背景和网格
        for plot in [self.plot_pitch, self.plot_roll, self.plot_yaw, self.plot_data]:
            # plot.setBackground('w')
            plot.showGrid(x=True, y=True)
            plot.addLegend()
        
        # 创建曲线
        colors = ['r', 'g', 'b', 'm']  # 红、绿、蓝、紫
        
        # Pitch图表曲线
        self.pitch_curves = []
        for i, color in enumerate(colors):
            self.pitch_curves.append(self.plot_pitch.plot(pen=color, name=f'IMU{i+1}'))
        
        # Roll图表曲线
        self.roll_curves = []
        for i, color in enumerate(colors):
            self.roll_curves.append(self.plot_roll.plot(pen=color, name=f'IMU{i+1}'))
        
        # Yaw图表曲线
        self.yaw_curves = []
        for i, color in enumerate(colors):
            self.yaw_curves.append(self.plot_yaw.plot(pen=color, name=f'IMU{i+1}'))
        
        # 附加数据图表曲线
        self.data_curves = []
        for i, color in enumerate(colors):
            self.data_curves.append(self.plot_data.plot(pen=color, name=f'Data{i+1}'))
    
    def update_plots(self, frame_data):
        """更新所有图表"""
        # 更新角度数据历史记录
        for i, imu in enumerate([frame_data.imu1, frame_data.imu2, frame_data.imu3, frame_data.imu4], 1):
            self.data_history[f'imu{i}_roll'].append(imu.roll)
            self.data_history[f'imu{i}_pitch'].append(imu.pitch)
            self.data_history[f'imu{i}_yaw'].append(imu.yaw)
        
        # 更新附加数据历史记录
        self.data_history['data1'].append(frame_data.data1)
        self.data_history['data2'].append(frame_data.data2)
        self.data_history['data3'].append(frame_data.data3)
        self.data_history['data4'].append(frame_data.data4)
        
        # 限制数据点数
        for key in self.data_history:
            if len(self.data_history[key]) > self.max_points:
                self.data_history[key] = self.data_history[key][-self.max_points:]
        
        # 准备X轴数据
        x_data = list(range(len(self.data_history['imu1_roll'])))
        
        # 更新Pitch图表
        for i, curve in enumerate(self.pitch_curves):
            curve.setData(x_data, self.data_history[f'imu{i+1}_pitch'])
        
        # 更新Roll图表
        for i, curve in enumerate(self.roll_curves):
            curve.setData(x_data, self.data_history[f'imu{i+1}_roll'])
        
        # 更新Yaw图表
        for i, curve in enumerate(self.yaw_curves):
            curve.setData(x_data, self.data_history[f'imu{i+1}_yaw'])
        
        # 更新附加数据图表
        for i, curve in enumerate(self.data_curves):
            curve.setData(x_data, self.data_history[f'data{i+1}'])
    
    def update_history_length(self, length):
        """更新历史数据长度"""
        self.max_points = length

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("IMU数据可视化系统")
        self.resize(1200, 900)
        
        self.serial_reader = SerialReader()
        self.data_saver = DataSaver()
        self.init_ui()
        self.init_serial_ports()
        
        # 定时器用于定期读取串口数据
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.serial_reader.read_data)
        self.timer.start(50)  # 20Hz更新
    
    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        main_layout = QVBoxLayout(central_widget)
        
        # 串口控制区域
        control_layout = QHBoxLayout()
        
        # 串口号选择
        self.port_combo = QComboBox()
        control_layout.addWidget(QLabel("串口:"))
        control_layout.addWidget(self.port_combo)
        
        # 波特率选择
        self.baudrate_combo = QComboBox()
        self.baudrate_combo.addItems(["9600", "19200", "38400", "57600", "115200", "230400", "460800", "921600"])
        self.baudrate_combo.setCurrentText("115200")
        control_layout.addWidget(QLabel("波特率:"))
        control_layout.addWidget(self.baudrate_combo)
        
        # 开始/停止按钮
        self.start_btn = QPushButton("开始采集")
        self.start_btn.clicked.connect(self.toggle_serial)
        control_layout.addWidget(self.start_btn)
        
        # 历史数据点数设置
        self.history_spin = QSpinBox()
        self.history_spin.setRange(100, 2000)
        self.history_spin.setValue(500)
        self.history_spin.setSuffix(" 点")
        self.history_spin.valueChanged.connect(self.update_history_length)
        
        control_layout.addWidget(QLabel("历史数据:"))
        control_layout.addWidget(self.history_spin)
        
        # 状态标签
        self.status_label = QLabel("就绪")
        control_layout.addWidget(self.status_label)
        
        main_layout.addLayout(control_layout)
        
        # 绘图区域
        self.plot_widget = MultiPlotWidget()
        main_layout.addWidget(self.plot_widget, 1)
        
        # 数据文本显示
        self.data_label = QLabel("等待数据...")
        self.data_label.setAlignment(Qt.AlignLeft)
        self.data_label.setStyleSheet("font-family: monospace;")
        main_layout.addWidget(self.data_label)
    
    def init_serial_ports(self):
        """获取可用串口列表"""
        self.port_combo.clear()
        ports = list_ports.comports()
        for port in ports:
            self.port_combo.addItem(port.device, port.device)
        if not ports:
            self.port_combo.addItem("无可用串口", "")
    
    def toggle_serial(self):
        """切换串口连接状态"""
        if self.serial_reader.serial_port and self.serial_reader.serial_port.is_open:
            # 停止采集时保存数据
            self.stop_collection()
        else:
            self.start_collection()
    
    def start_collection(self):
        """开始数据采集"""
        port = self.port_combo.currentData()
        baudrate = int(self.baudrate_combo.currentText())
        if self.serial_reader.start(port, baudrate, self.process_frame_data):
            self.start_btn.setText("停止采集")
            self.status_label.setText("采集数据中...")
            self.data_label.setText("等待数据...")
            self.data_saver.clear()  # 清空之前的记录
        else:
            self.status_label.setText("串口打开失败")
    
    def stop_collection(self):
        """停止数据采集并保存数据"""
        self.serial_reader.stop()
        self.start_btn.setText("开始采集")
        self.status_label.setText("就绪")
        
        # 保存数据
        if self.data_saver.records:
            options = QFileDialog.Options()
            filename, _ = QFileDialog.getSaveFileName(
                self, "保存数据", "", "CSV文件 (*.csv);;所有文件 (*)", options=options)
            
            if filename:
                if not filename.endswith('.csv'):
                    filename += '.csv'
                
                if self.data_saver.save_to_csv(filename):
                    QMessageBox.information(self, "保存成功", f"数据已保存到 {filename}")
                else:
                    QMessageBox.warning(self, "保存失败", "数据保存失败")
    
    def process_frame_data(self, frame_data):
        """处理接收到的帧数据"""
        # 更新图表
        self.plot_widget.update_plots(frame_data)
        
        # 保存数据
        self.data_saver.add_record(frame_data)
        
        # 更新文本显示(显示第一个IMU的数据作为示例)
        imu = frame_data.imu1
        text = f"""
        IMU1 最新数据:
        俯仰: {imu.pitch:.2f}°
        横滚: {imu.roll:.2f}°
        偏航: {imu.yaw:.2f}°
        温度: {imu.temperature}
        
        附加数据:
        Data1: {frame_data.data1}
        Data2: {frame_data.data2}
        Data3: {frame_data.data3}
        Data4: {frame_data.data4}
        
        已采集帧数: {len(self.data_saver.records)}
        """
        self.data_label.setText(text)
    
    def update_history_length(self, length):
        """更新历史数据长度"""
        self.plot_widget.update_history_length(length)
    
    def closeEvent(self, event):
        """窗口关闭时停止串口"""
        self.serial_reader.stop()
        event.accept()

class SerialReader:
    def __init__(self):
        self.serial_port = None
        self.buffer = bytearray()
        self.callback = None
    
    def start(self, port, baudrate, callback):
        """开始读取串口数据"""
        self.callback = callback
        try:
            self.serial_port = serial.Serial(port, baudrate, timeout=0.1)
            return True
        except Exception as e:
            print(f"串口打开失败: {e}")
            return False
    
    def stop(self):
        """停止读取串口数据"""
        if self.serial_port and self.serial_port.is_open:
            self.serial_port.close()
        self.serial_port = None
    
    def read_data(self):
        """读取并处理串口数据"""
        if not self.serial_port or not self.serial_port.is_open:
            return
        
        data = self.serial_port.read(self.serial_port.in_waiting or 1)
        self.buffer.extend(data)
        
        while len(self.buffer) >= 66:
            header_pos = self.buffer.find(b'\xA5\x5A')
            if header_pos == -1:
                self.buffer.clear()
                break
            
            if header_pos > 0:
                del self.buffer[:header_pos]
            
            if len(self.buffer) < 66:
                break
            
            frame_data = self.buffer[:66]
            del self.buffer[:66]
            
            try:
                parsed = IMUParser.parse_frame(frame_data)
                if self.callback:
                    self.callback(parsed)
            except Exception as e:
                print(f"解析错误: {e}")
                continue

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())